"""
可视化模块
提供股票数据的各种图表绘制功能
"""

import matplotlib.pyplot as plt
import pandas as pd
import numpy as np
import streamlit as st
import plotly.graph_objects as go
from plotly.subplots import make_subplots
import matplotlib.dates as mdates


# 设置中文字体
plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False


def plot_stock_price(df, title="茅台股票价格走势", height=600):
    """
    绘制股票价格走势图（使用Plotly）
    
    Args:
        df: 股票数据
        title: 图表标题
        height: 图表高度
    """
    fig = go.Figure()
    
    # 添加收盘价线
    fig.add_trace(go.Scatter(
        x=df.index,
        y=df['close'],
        mode='lines',
        name='收盘价',
        line=dict(color='red', width=2)
    ))
    
    # 添加移动平均线
    if 'ma5' in df.columns:
        fig.add_trace(go.Scatter(
            x=df.index,
            y=df['ma5'],
            mode='lines',
            name='MA5',
            line=dict(color='blue', width=1),
            opacity=0.7
        ))
    
    if 'ma20' in df.columns:
        fig.add_trace(go.Scatter(
            x=df.index,
            y=df['ma20'],
            mode='lines',
            name='MA20',
            line=dict(color='green', width=1),
            opacity=0.7
        ))
    
    fig.update_layout(
        title=title,
        xaxis_title="日期",
        yaxis_title="价格（元）",
        height=height,
        showlegend=True,
        hovermode='x unified'
    )
    
    return fig


def plot_candlestick(df, title="茅台股票K线图", height=600, show_volume=True):
    """
    绘制K线图
    
    Args:
        df: 股票数据
        title: 图表标题
        height: 图表高度
        show_volume: 是否显示成交量
    """
    if show_volume:
        fig = make_subplots(
            rows=2, cols=1,
            shared_xaxes=True,
            vertical_spacing=0.1,
            subplot_titles=('K线图', '成交量'),
            row_width=[0.2, 0.7]
        )
    else:
        fig = go.Figure()
    
    # K线图
    candlestick = go.Candlestick(
        x=df.index,
        open=df['open'],
        high=df['high'],
        low=df['low'],
        close=df['close'],
        name="K线"
    )
    
    if show_volume:
        fig.add_trace(candlestick, row=1, col=1)
        
        # 成交量
        colors = ['red' if close >= open else 'green' 
                 for close, open in zip(df['close'], df['open'])]
        
        fig.add_trace(go.Bar(
            x=df.index,
            y=df['volume'],
            name="成交量",
            marker_color=colors,
            opacity=0.7
        ), row=2, col=1)
        
    else:
        fig.add_trace(candlestick)
    
    fig.update_layout(
        title=title,
        height=height,
        xaxis_rangeslider_visible=False,
        showlegend=True
    )
    
    return fig


def plot_technical_indicators(df, height=800):
    """
    绘制技术指标图表
    
    Args:
        df: 包含技术指标的股票数据
        height: 图表高度
    """
    fig = make_subplots(
        rows=3, cols=1,
        shared_xaxes=True,
        vertical_spacing=0.08,
        subplot_titles=('价格与移动平均线', 'RSI', 'MACD'),
        row_heights=[0.5, 0.25, 0.25]
    )
    
    # 第一行：价格和移动平均线
    fig.add_trace(go.Scatter(
        x=df.index, y=df['close'],
        mode='lines', name='收盘价',
        line=dict(color='black', width=2)
    ), row=1, col=1)
    
    if 'ma5' in df.columns:
        fig.add_trace(go.Scatter(
            x=df.index, y=df['ma5'],
            mode='lines', name='MA5',
            line=dict(color='blue', width=1)
        ), row=1, col=1)
    
    if 'ma20' in df.columns:
        fig.add_trace(go.Scatter(
            x=df.index, y=df['ma20'],
            mode='lines', name='MA20',
            line=dict(color='red', width=1)
        ), row=1, col=1)
    
    # 布林带
    if 'bollinger_upper' in df.columns:
        fig.add_trace(go.Scatter(
            x=df.index, y=df['bollinger_upper'],
            mode='lines', name='布林上轨',
            line=dict(color='gray', width=1, dash='dash'),
            opacity=0.5
        ), row=1, col=1)
        
        fig.add_trace(go.Scatter(
            x=df.index, y=df['bollinger_lower'],
            mode='lines', name='布林下轨',
            line=dict(color='gray', width=1, dash='dash'),
            opacity=0.5,
            fill='tonexty', fillcolor='rgba(128,128,128,0.1)'
        ), row=1, col=1)
    
    # 第二行：RSI
    if 'rsi' in df.columns:
        fig.add_trace(go.Scatter(
            x=df.index, y=df['rsi'],
            mode='lines', name='RSI',
            line=dict(color='purple', width=2)
        ), row=2, col=1)
        
        # RSI超买超卖线
        fig.add_hline(y=70, line_dash="dash", line_color="red", 
                     annotation_text="超买", row=2, col=1)
        fig.add_hline(y=30, line_dash="dash", line_color="green", 
                     annotation_text="超卖", row=2, col=1)
    
    # 第三行：MACD
    if 'macd' in df.columns:
        fig.add_trace(go.Scatter(
            x=df.index, y=df['macd'],
            mode='lines', name='MACD',
            line=dict(color='blue', width=2)
        ), row=3, col=1)
        
        fig.add_trace(go.Scatter(
            x=df.index, y=df['macd_signal'],
            mode='lines', name='MACD信号',
            line=dict(color='red', width=1)
        ), row=3, col=1)
        
        # MACD柱状图
        colors = ['red' if val >= 0 else 'green' for val in df['macd_histogram']]
        fig.add_trace(go.Bar(
            x=df.index, y=df['macd_histogram'],
            name='MACD柱', marker_color=colors,
            opacity=0.6
        ), row=3, col=1)
    
    fig.update_layout(
        height=height,
        title="技术指标分析",
        showlegend=True,
        hovermode='x unified'
    )
    
    # 更新y轴标题
    fig.update_yaxes(title_text="价格（元）", row=1, col=1)
    fig.update_yaxes(title_text="RSI", row=2, col=1)
    fig.update_yaxes(title_text="MACD", row=3, col=1)
    
    return fig


def plot_prediction_results(df, predictions, title="股价预测结果"):
    """
    绘制预测结果图表
    
    Args:
        df: 历史股票数据
        predictions: 预测结果字典
        title: 图表标题
    """
    fig = go.Figure()
    
    # 显示最近的历史数据（60天）
    recent_days = min(60, len(df))
    recent_df = df.tail(recent_days)
    
    # 历史价格
    fig.add_trace(go.Scatter(
        x=recent_df.index,
        y=recent_df['close'],
        mode='lines',
        name='历史价格',
        line=dict(color='blue', width=2)
    ))
    
    # 预测价格
    fig.add_trace(go.Scatter(
        x=predictions['future_dates'],
        y=predictions['future_prices'],
        mode='lines+markers',
        name='预测价格',
        line=dict(color='red', width=2),
        marker=dict(size=6)
    ))
    
    # 连接线
    last_date = recent_df.index[-1]
    last_price = recent_df['close'].iloc[-1]
    first_pred_date = predictions['future_dates'][0]
    first_pred_price = predictions['future_prices'][0]
    
    fig.add_trace(go.Scatter(
        x=[last_date, first_pred_date],
        y=[last_price, first_pred_price],
        mode='lines',
        name='连接线',
        line=dict(color='orange', width=2, dash='dash'),
        opacity=0.7
    ))
    
    fig.update_layout(
        title=title,
        xaxis_title="日期",
        yaxis_title="股价（元）",
        height=600,
        showlegend=True,
        hovermode='x unified'
    )
    
    return fig


def plot_model_performance(results):
    """
    绘制模型性能图表
    
    Args:
        results: 模型训练结果
    """
    fig = go.Figure()
    
    # 真实价格
    fig.add_trace(go.Scatter(
        x=results['test_dates'],
        y=results['y_test_real'],
        mode='lines',
        name='真实价格',
        line=dict(color='blue', width=2)
    ))
    
    # 预测价格
    fig.add_trace(go.Scatter(
        x=results['test_dates'],
        y=results['y_pred_test_real'],
        mode='lines',
        name='预测价格',
        line=dict(color='red', width=2),
        opacity=0.8
    ))
    
    fig.update_layout(
        title="模型预测性能",
        xaxis_title="日期",
        yaxis_title="股价（元）",
        height=500,
        showlegend=True,
        hovermode='x unified'
    )
    
    return fig


def display_metrics(metrics, title="模型评估指标"):
    """
    显示评估指标
    
    Args:
        metrics: 指标字典
        title: 标题
    """
    st.subheader(title)
    
    col1, col2, col3 = st.columns(3)
    
    with col1:
        st.metric(
            label="均方误差 (MSE)",
            value=f"{metrics['mse']:.2f}"
        )
    
    with col2:
        st.metric(
            label="平均绝对误差 (MAE)",
            value=f"{metrics['mae']:.2f}"
        )
    
    with col3:
        st.metric(
            label="R² 得分",
            value=f"{metrics['r2']:.4f}",
            delta=f"{(metrics['r2']*100):.2f}% 准确率"
        )